#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MongoDB数据备份工具
在进行数据库优化前备份现有数据
"""

import json
import os
from datetime import datetime
from pymongo import MongoClient
from typing import List, Dict, Any

class MongoDBBackup:
    """MongoDB数据备份类"""
    
    def __init__(self, connection_string: str = "mongodb://localhost:27017/", db_name: str = "civil_code_db"):
        self.connection_string = connection_string
        self.db_name = db_name
        self.backup_dir = "backup"
        
    def create_backup(self) -> bool:
        """创建完整数据备份"""
        print("📦 开始创建MongoDB数据备份...")
        
        try:
            # 创建备份目录
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            backup_path = os.path.join(self.backup_dir, f"civil_code_backup_{timestamp}")
            os.makedirs(backup_path, exist_ok=True)
            
            # 连接数据库
            client = MongoClient(self.connection_string)
            db = client[self.db_name]
            
            # 备份主要集合
            collections_to_backup = [
                'civil_code_articles',
                'user_favorites', 
                'search_history'
            ]
            
            backup_info = {
                "backup_time": datetime.now().isoformat(),
                "database": self.db_name,
                "collections": {}
            }
            
            for collection_name in collections_to_backup:
                if collection_name in db.list_collection_names():
                    collection = db[collection_name]
                    
                    # 导出数据
                    documents = list(collection.find({}))
                    count = len(documents)
                    
                    # 处理ObjectId序列化
                    for doc in documents:
                        if '_id' in doc:
                            doc['_id'] = str(doc['_id'])
                    
                    # 保存到文件
                    backup_file = os.path.join(backup_path, f"{collection_name}.json")
                    with open(backup_file, 'w', encoding='utf-8') as f:
                        json.dump(documents, f, ensure_ascii=False, indent=2, default=str)
                    
                    backup_info["collections"][collection_name] = {
                        "document_count": count,
                        "backup_file": f"{collection_name}.json"
                    }
                    
                    print(f"   ✅ 备份集合 {collection_name}: {count} 条记录")
                else:
                    print(f"   ⚠️ 集合 {collection_name} 不存在，跳过备份")
            
            # 备份索引信息
            indexes_info = {}
            for collection_name in collections_to_backup:
                if collection_name in db.list_collection_names():
                    collection = db[collection_name]
                    indexes = list(collection.list_indexes())
                    indexes_info[collection_name] = []
                    
                    for index in indexes:
                        # 转换为可序列化的格式
                        index_info = {
                            "name": index.get("name"),
                            "key": dict(index.get("key", {})),
                            "unique": index.get("unique", False),
                            "sparse": index.get("sparse", False)
                        }
                        indexes_info[collection_name].append(index_info)
            
            # 保存备份信息和索引信息
            with open(os.path.join(backup_path, "backup_info.json"), 'w', encoding='utf-8') as f:
                json.dump(backup_info, f, ensure_ascii=False, indent=2, default=str)
            
            with open(os.path.join(backup_path, "indexes_info.json"), 'w', encoding='utf-8') as f:
                json.dump(indexes_info, f, ensure_ascii=False, indent=2, default=str)
            
            print(f"✅ 数据备份完成！")
            print(f"📁 备份路径: {backup_path}")
            print(f"📊 备份统计: {sum(info['document_count'] for info in backup_info['collections'].values())} 条记录")
            
            client.close()
            return True
            
        except Exception as e:
            print(f"❌ 备份失败: {str(e)}")
            return False
    
    def verify_backup(self, backup_path: str) -> bool:
        """验证备份完整性"""
        print(f"🔍 验证备份完整性: {backup_path}")
        
        try:
            # 检查备份信息文件
            backup_info_file = os.path.join(backup_path, "backup_info.json")
            if not os.path.exists(backup_info_file):
                print("❌ 备份信息文件不存在")
                return False
            
            with open(backup_info_file, 'r', encoding='utf-8') as f:
                backup_info = json.load(f)
            
            # 验证每个集合的备份文件
            for collection_name, info in backup_info["collections"].items():
                backup_file = os.path.join(backup_path, info["backup_file"])
                if not os.path.exists(backup_file):
                    print(f"❌ 集合 {collection_name} 的备份文件不存在")
                    return False
                
                # 验证文件内容
                with open(backup_file, 'r', encoding='utf-8') as f:
                    documents = json.load(f)
                
                if len(documents) != info["document_count"]:
                    print(f"❌ 集合 {collection_name} 记录数不匹配")
                    return False
                
                print(f"   ✅ 集合 {collection_name}: {len(documents)} 条记录")
            
            print("✅ 备份验证通过")
            return True
            
        except Exception as e:
            print(f"❌ 备份验证失败: {str(e)}")
            return False

def main():
    """主函数"""
    print("=" * 60)
    print("📦 MongoDB数据备份工具")
    print("=" * 60)
    
    backup_tool = MongoDBBackup()
    
    # 创建备份
    if backup_tool.create_backup():
        print("\n🎉 数据备份成功完成！")
    else:
        print("\n❌ 数据备份失败")

if __name__ == "__main__":
    main()
